import os
import torch
# from PIL import Image
from sklearn.decomposition import PCA
from sklearn.manifold import TSNE
from tqdm import tqdm
import argparse
import matplotlib.pyplot as plt
from utils.common import check_and_create_directory, list_all_files, chunk_list

if __name__ == '__main__':
    # todo 提取所有的数据，然后合并在一起，看看大小
    input_path = '/home/zry/datasets/building/train/rgbs'
    output_path = '/home/zry/datasets/building/train/seg'
    tmp_path = '/home/zry/datasets/building/train/tmp'

    # 所有pth
    paths = list_all_files(tmp_path, '.pth')
    encs = []
    for path in tqdm(paths):
        f = torch.load(path, "cpu")
        idx, seg_images, enc = f['idx'], f['seg_images'], f['enc']
        encs.append(enc)
        # break
        
    encs = torch.cat(encs)
    encs = encs.detach().numpy()
    pca = PCA(n_components=2)
    
    reduced_data = pca.fit_transform(encs)
    plt.figure(figsize=(100, 60))
    plt.scatter(reduced_data[:, 0], reduced_data[:, 1], c='blue', marker='o')
    plt.grid(True)

    # 保存图片文件
    plt.savefig('pca_reduction.png')



